package com.sinosig.slbi.gm;

import org.bouncycastle.asn1.ASN1OctetString;
import org.bouncycastle.asn1.DERBMPString;
import org.bouncycastle.asn1.DEROctetString;
import org.bouncycastle.asn1.pkcs.PKCSObjectIdentifiers;
import org.bouncycastle.asn1.x500.X500Name;
import org.bouncycastle.asn1.x509.Extension;
import org.bouncycastle.asn1.x509.KeyUsage;
import org.bouncycastle.cert.X509CertificateHolder;
import org.bouncycastle.cert.jcajce.JcaX509CertificateConverter;
import org.bouncycastle.cert.jcajce.JcaX509v3CertificateBuilder;
import org.bouncycastle.crypto.engines.DESedeEngine;
import org.bouncycastle.crypto.modes.CBCBlockCipher;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.operator.ContentSigner;
import org.bouncycastle.operator.jcajce.JcaContentSignerBuilder;
import org.bouncycastle.pkcs.PKCS12PfxPdu;
import org.bouncycastle.pkcs.PKCS12PfxPduBuilder;
import org.bouncycastle.pkcs.PKCS12SafeBag;
import org.bouncycastle.pkcs.PKCS12SafeBagBuilder;
import org.bouncycastle.pkcs.bc.BcPKCS12MacCalculatorBuilder;
import org.bouncycastle.pkcs.bc.BcPKCS12PBEOutputEncryptorBuilder;
import org.bouncycastle.pkcs.jcajce.JcaPKCS12SafeBagBuilder;

import java.io.FileOutputStream;
import java.math.BigInteger;
import java.security.*;
import java.security.cert.X509Certificate;
import java.security.spec.ECGenParameterSpec;
import java.util.Date;

public class SM2CertGenerator {

    static {
        Security.addProvider(new BouncyCastleProvider());
    }

    public static void main(String[] args) throws Exception {

//        String cerName = "sign.cer";
//        String pfxName = "sign.pfx";
//        String password = "123456";
//        String flag = "sign";

        String cerName = "enc.cer";
        String pfxName = "enc.pfx";
        String password = "123456";
        String flag = "enc";
        // 生成密钥对
        KeyPair keyPair = generateSM2KeyPair();

        // 生成X509证书
        X509Certificate certificate = generateCertificate(keyPair,flag);

        // 保存公钥为cer格式
        savePublicKeyToCer(certificate, cerName);

        // 保存私钥为pfx格式
        savePrivateKeyToPfx(keyPair.getPrivate(), certificate, pfxName, password);
    }

    private static KeyPair generateSM2KeyPair() throws Exception {
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("EC", "BC");
        keyPairGenerator.initialize(new ECGenParameterSpec("sm2p256v1"));
        return keyPairGenerator.generateKeyPair();
    }

    private static X509Certificate generateCertificate(KeyPair keyPair,String flag) throws Exception {
        X500Name issuer = new X500Name("CN=SM2 Cert, O=My Organization, C=CN");
        X500Name subject = issuer;
        BigInteger serial = BigInteger.valueOf(System.currentTimeMillis());
        Date notBefore = new Date(System.currentTimeMillis() - 24 * 60 * 60 * 1000);
        Date notAfter = new Date(System.currentTimeMillis() + 365L * 24 * 60 * 60 * 10000);

        JcaX509v3CertificateBuilder certBuilder = new JcaX509v3CertificateBuilder(
                issuer,
                serial,
                notBefore,
                notAfter,
                subject,
                keyPair.getPublic()
        );

        if ("sign".equals(flag)){
            // 添加密钥用途扩展（关键扩展）
            KeyUsage keyUsage = new KeyUsage(
                            KeyUsage.digitalSignature |
                            KeyUsage.nonRepudiation
            );
            certBuilder.addExtension(
                    Extension.keyUsage,
                    true,  // 是否关键扩展（true表示必须被校验）
                    keyUsage
            );
        }
        if ("enc".equals(flag)){
            // 添加密钥用途扩展（关键扩展）
            KeyUsage keyUsage = new KeyUsage(
                            KeyUsage.keyEncipherment |
                            KeyUsage.dataEncipherment|
                            KeyUsage.keyAgreement
            );
            certBuilder.addExtension(
                    Extension.keyUsage,
                    true,  // 是否关键扩展（true表示必须被校验）
                    keyUsage
            );
        }

        ContentSigner signer = new JcaContentSignerBuilder("SM3withSM2")
                .setProvider("BC")
                .build(keyPair.getPrivate());

        X509CertificateHolder certHolder = certBuilder.build(signer);
        return new JcaX509CertificateConverter()
                .setProvider("BC")
                .getCertificate(certHolder);
    }

    private static void savePublicKeyToCer(X509Certificate certificate, String filename) throws Exception {
        try (FileOutputStream fos = new FileOutputStream(filename)) {
            fos.write(certificate.getEncoded());
        }
    }

    private static void savePrivateKeyToPfx(PrivateKey privateKey, 
                                          X509Certificate certificate,
                                          String filename, 
                                          String password) throws Exception {
        PKCS12PfxPduBuilder pfxBuilder = new PKCS12PfxPduBuilder();
        // 生成密钥本地标识（必须与证书标识一致）
        byte[] keyId = new byte[20];
        new SecureRandom().nextBytes(keyId);
        ASN1OctetString keyIdentifier = new DEROctetString(keyId);
        // 创建密钥安全包
        PKCS12SafeBagBuilder keyBagBuilder = new JcaPKCS12SafeBagBuilder(privateKey)
                .addBagAttribute(PKCSObjectIdentifiers.pkcs_9_at_friendlyName, new DERBMPString("SM2"))
                .addBagAttribute(PKCSObjectIdentifiers.pkcs_9_at_localKeyId, keyIdentifier); // 关键属性
        // 创建证书安全包
        PKCS12SafeBagBuilder certBagBuilder = new JcaPKCS12SafeBagBuilder(certificate)
                .addBagAttribute(PKCSObjectIdentifiers.pkcs_9_at_friendlyName,
                        new DERBMPString("SM2"));

        PKCS12SafeBag[] pkcs12SafeBags = {keyBagBuilder.build(),certBagBuilder.build()};
        // 构建PFX结构
        pfxBuilder.addEncryptedData(
                new BcPKCS12PBEOutputEncryptorBuilder(
                        PKCSObjectIdentifiers.pbeWithSHAAnd3_KeyTripleDES_CBC,
                        new CBCBlockCipher(new DESedeEngine())).build(password.toCharArray()),
                pkcs12SafeBags);

        PKCS12PfxPdu pfx = pfxBuilder.build(
                new BcPKCS12MacCalculatorBuilder(),
                password.toCharArray());

        try (FileOutputStream fos = new FileOutputStream(filename)) {
            fos.write(pfx.getEncoded());
        }
    }
}